import bpmn.structure.*;
import bpmn.structure.Process;
import common.NameTagUtils;
import ta.TALayoutHelper;
import ta.structure.TADefinition;
import ta.structure.TALocation;
import ta.structure.TATemplate;
import ta.structure.TATransition;

import java.util.*;

public class BPMN2TATransformer
{
    public Collection<BPMNDefinition> bpmns;
    public TALayoutHelper layoutHelper;
    public TADefinition ta;
    public NameTagUtils ng;
    // bpmn lane id --> current lane template
    private Map<String, TATemplate> templateSet;

    // need reset after single transform
    private final Stack<BPMNElemType> stateStack;
    private final HashMap<String, TATemplate> templateRefMap;
    private final HashMap<String, BPMNDefinition> nameBPMNMap;
    private final HashMap<String, Integer> outRefMap;
    private final Stack<HashMap<String, TALocation>> branchStateStack;
    private final HashMap<String, TALocation> nodeMap;
    private final HashSet<String> nodeFlag;
    private final List<Map<String, String>> callLinks;

    public BPMN2TATransformer(Collection<BPMNDefinition> defs)
    {
        bpmns = defs;
        ta = new TADefinition();
        ng = new NameTagUtils();
        // outgoing ref num map
        nodeMap = new HashMap<>();
        outRefMap = new HashMap<>();
        templateRefMap = new HashMap<>();
        nameBPMNMap = new HashMap<>();
        stateStack = new Stack<>();
        branchStateStack = new Stack<>();
        callLinks = new ArrayList<>();
        layoutHelper = new TALayoutHelper();
        nodeFlag = new HashSet<>();
        reset();
    }

    public void transform(String ta_name)
    {
        // TA model name, also file name;
        ta.setName(ta_name);

        // read bpmn name
        for (var bpmn : bpmns)
            nameBPMNMap.put(bpmn.getFileName(), bpmn);

        // single transform
        for (var bpmn : bpmns)
        {
            reset();
            layoutHelper.reset();
            transformSingle(bpmn);
        }

        // linking
        for (var map : callLinks)
        {
            String called = map.get("bpmn_name");
            String startSig = map.get("start");
            String endSig = map.get("end");
            TATemplate targetT = nameBPMNMap.get(called).getMainTemplate();
            String base = targetT.getName() + "_MAIN";
            if (startSig != null)
            {
                for (var tran : targetT.getStartTrans())
                {
                    tran.setSync(base + "_START", false);
                }
                targetT.addParameter("chan&", base + "_START");
                targetT.setStartSignal(startSig);
            }
            if (endSig != null)
            {
                for (var tran : targetT.getFinishTrans())
                {
                    tran.setSync(base + "_END", true);
                }
                targetT.addParameter("chan&", base + "_END");
                targetT.setFinishSignal(endSig);
            }
        }
        return;
    }

    private void reset()
    {
        stateStack.clear();
        templateRefMap.clear();
        outRefMap.clear();
        branchStateStack.clear();
        nodeMap.clear();
        nodeFlag.clear();
        layoutHelper = new TALayoutHelper();
    }

    private void transformSingle(BPMNDefinition bpmn)
    {
        var process = bpmn.getProcessList().get(0);

        if (process.getLaneSet() == null)
        {
            // template setup
            TATemplate t = new TATemplate(process.id);
            String name = process.getName();
            t.setName(name);
            ta.addTemplate(t);

            // init location setup
            TALocation srt = newInitLocation(name, name, t);

            layoutHelper.reset();
            layoutHelper.layoutLocation(srt);
            outRefMap.put(srt.getId(), 0);

            t.setInitRef(srt.getId());
            t.addLocation(srt);

            for (var node : process.getElements())
            {
                templateRefMap.put(node.id, t);
            }
            bpmn.setMainTemplate(t);
        }
        else
        {
            layoutHelper.reset();
            TATemplate ct = generateControlTemplate(process);
            ta.addTemplate(ct);
            bpmn.setMainTemplate(ct);
            templateSet = generateTemplateSet(process);
        }

        ProcessElem start = null;
        for (var e : process.getElements())
        {
            if (e instanceof StartEvent)
            {
                start = e;
                break;
            }
        }
        if (start == null)
            errorLog("no StartEvent found in bpmn");
        assert start != null;

        DFSTrans(bpmn, start, templateRefMap.get(start.id).getPtr(), null);
    }

    private void DFSTrans(BPMNDefinition bpmn, ProcessElem elem, TALocation prev, SequenceFlow flowRef)
    {
        BPMNElemType t = elem.type;
        switch (elem.type)
        {
            case StartEvent:
            {
                TALocation loc = linkLocation(elem, prev, flowRef, false);
                TATemplate tt = loc.getParent();
                TATransition srt = tt.getTransitions().get(0);
                String cName = bpmn.getProcessList().get(0).getName() + "_MAIN";
                srt.setSync(cName + "_START", false);
                nodeMap.put(loc.getId(), loc);
                tt.addStartTrans(srt);

                // get next node
                String nxtFlowRef = elem.getOutgoingRefs().get(0);
                SequenceFlow nxtFlow = (SequenceFlow) bpmn.elemMap.get(nxtFlowRef);
                ProcessElem nxt = (ProcessElem) bpmn.elemMap.get(nxtFlow.getTargetRef());
                if (nodeFlag.contains(nxt.id))
                {
                    setTransitionNoLayout(loc, nodeMap.get(ng.queryLastIdHitFix(nxt.id)), nxtFlow);
                }
                else
                {
                    nodeFlag.add(nxt.id);
                    // dive
                    stateStack.push(t);
                    DFSTrans(bpmn, nxt, loc, nxtFlow);
                    stateStack.pop();
                }
                break;
            }
            case Task:
            {
                TATemplate fromT = prev.getParent();
                TATemplate toT = templateRefMap.get(elem.id);

                TALocation loc = linkLocation(elem, prev, flowRef, !fromT.equals(toT));
                nodeMap.put(loc.getId(), loc);

                // get next node
                String nxtFlowRef = elem.getOutgoingRefs().get(0);
                SequenceFlow nxtFlow = (SequenceFlow) bpmn.elemMap.get(nxtFlowRef);
                ProcessElem nxt = (ProcessElem) bpmn.elemMap.get(nxtFlow.getTargetRef());
                if (nodeFlag.contains(nxt.id))
                {
                    if (!(nxt instanceof ExclusiveGateway))
                        setTransitionNoLayout(loc, nodeMap.get(ng.queryLastIdHitFix(nxt.id)), nxtFlow);
                    else
                    {
                        stateStack.push(t);
                        DFSTrans(bpmn, nxt, loc, nxtFlow);
                        stateStack.pop();
                    }
                }
                else
                {
                    nodeFlag.add(nxt.id);
                    // dive
                    stateStack.push(t);
                    DFSTrans(bpmn, nxt, loc, nxtFlow);
                    stateStack.pop();
                }
                break;
            }
            case SubProcess:
            {
                if (!nameBPMNMap.containsKey(elem.getName()))
                    errorLog("Missing bpmn diagram: " + elem.getName() + ".bpmn");
                TALocation cloc = linkLocationSub(elem, prev, flowRef);

                String nameId = String.format("%s_%s",
                                              ng.queryLastNameHitFix(elem.getName()),
                                              ng.queryLastIdHitFix(elem.id));
                HashMap<String, String> map = new HashMap<>();
                map.put("bpmn_name", elem.getName());
                map.put("start", nameId + "_START");
                map.put("end", nameId + "_END");
                callLinks.add(map);

                // get next node
                String nxtFlowRef = elem.getOutgoingRefs().get(0);
                SequenceFlow nxtFlow = (SequenceFlow) bpmn.elemMap.get(nxtFlowRef);
                ProcessElem nxt = (ProcessElem) bpmn.elemMap.get(nxtFlow.getTargetRef());
                if (nodeFlag.contains(nxt.id))
                {
                    if (!(nxt instanceof ExclusiveGateway))
                        setTransitionNoLayout(cloc, nodeMap.get(ng.queryLastIdHitFix(nxt.id)), nxtFlow);
                    else
                    {
                        stateStack.push(t);
                        DFSTrans(bpmn, nxt, cloc, nxtFlow);
                        stateStack.pop();
                    }
                }
                else
                {
                    nodeFlag.add(nxt.id);
                    // dive
                    stateStack.push(t);
                    DFSTrans(bpmn, nxt, cloc, nxtFlow);
                    stateStack.pop();
                }
                break;
            }
            case EndEvent:
            {
                TATemplate fromT = prev.getParent();
                TATemplate toT = templateRefMap.get(elem.id);
                Process p = bpmn.getProcessList().get(0);

                TALocation loc = linkLocation(elem, prev, flowRef, !fromT.equals(toT));
                TALocation endLoc = newFinishLocation(p.id, p.getName(), toT);
                SequenceFlow seq = new SequenceFlow(loc.getId() + "_END");
                seq.setName(loc.getName() + "_END");

                String cName = p.getName() + "_MAIN";
                TATransition trans = setTransitionPoolSame(loc, endLoc, seq);
                toT.addLocation(endLoc);
                trans.setSync(cName + "_END", true);
                toT.addFinishTrans(trans);
                return;
            }
            case ExclusiveGateway:
            {
                ExclusiveGateway eg = (ExclusiveGateway) elem;
                // save state
                HashMap<String, TALocation> map = new HashMap<>();
                for (var tp : ta.getTemplates().values())
                {
                    map.put(tp.getId(), tp.getPtr());
                }
                branchStateStack.push(map);
                stateStack.push(t);
                for (int i = 0; i < eg.getOutgoingRefs().size(); i++)
                {
                    // get next node in this branch
                    SequenceFlow nxtFlow = (SequenceFlow) bpmn.elemMap.get(eg.getOutgoingRefs().get(i));
                    ProcessElem nxt = (ProcessElem) bpmn.elemMap.get(nxtFlow.getTargetRef());

                    if (nodeFlag.contains(nxt.id))
                    {
                        if (!(nxt instanceof ExclusiveGateway))
                            setTransitionNoLayout(prev, nodeMap.get(ng.queryLastIdHitFix(nxt.id)), nxtFlow);
                        else
                        {
                            stateStack.push(t);
                            DFSTrans(bpmn, nxt, prev, nxtFlow);
                            stateStack.pop();
                        }
                    }
                    else
                    {
                        nodeFlag.add(nxt.id);
                        // dive
                        stateStack.push(t);
                        DFSTrans(bpmn, nxt, prev, nxtFlow);
                        stateStack.pop();
                    }
                }
                stateStack.pop();
                branchStateStack.pop();
                break;
            }
            case ParallelGateway:
            {
                ParallelGateway pg = (ParallelGateway) elem;
                // generate new template set
                TATemplate tt = templateRefMap.get(pg.id);
                TALocation endLoc = newSignalLocation(pg.id, pg.getName(), tt);
                TATransition trans = newTransition(pg.id, pg.getName(), prev.getId(), endLoc
                        .getId(), "_PARALLEL_SIGNAL");
                TATransition cycle = newTransition(pg.id, pg.getName(), endLoc.getId(), endLoc
                        .getId(), "_PARALLEL_CYCLE");
                tt.addLocation(endLoc);
                tt.addTransition(trans);
                tt.addTransition(cycle);
                layoutHelper.layoutLocation(prev, endLoc);
                layoutHelper.layoutTransition(trans, prev, endLoc);
                layoutHelper.layoutTransition(cycle, endLoc, endLoc);
                trans.setSync(String.format("%s_%s_PARALLEL", pg.getName(), pg.id), true);
                cycle.setSync(String.format("%s_%s_PARALLEL", pg.getName(), pg.id), true);
                ta.getGlobals().addVariable("chan", String.format("%s_%s_PARALLEL", pg.getName(), pg.id));

                for (int i = 0; i < pg.getOutgoingRefs().size(); i++)
                {
                    templateSet = generateTemplateSet(bpmn.getProcessList().get(0));
                    // get next node in this branch
                    SequenceFlow nxtFlow = (SequenceFlow) bpmn.elemMap.get(pg.getOutgoingRefs().get(i));
                    ProcessElem nxt = (ProcessElem) bpmn.elemMap.get(nxtFlow.getTargetRef());

                    if (nodeFlag.contains(nxt.id))
                    {
                        if (!(nxt instanceof ExclusiveGateway))
                            setTransitionNoLayout(prev, nodeMap.get(ng.queryLastIdHitFix(nxt.id)), nxtFlow);
                        else
                        {
                            stateStack.push(t);
                            DFSTrans(bpmn, nxt, templateRefMap.get(nxt.id).getPtr(), nxtFlow);
                            templateRefMap.get(nxt.id).getTransitions().get(0)
                                    .setSync(String.format("%s_%s_PARALLEL", pg.getName(), pg.id), false);
                            stateStack.pop();
                        }
                    }
                    else
                    {
                        nodeFlag.add(nxt.id);
                        // dive
                        stateStack.push(t);
                        DFSTrans(bpmn, nxt, templateRefMap.get(nxt.id).getPtr(), nxtFlow);
                        templateRefMap.get(nxt.id).getTransitions().get(0)
                                .setSync(String.format("%s_%s_PARALLEL", pg.getName(), pg.id), false);
                        stateStack.pop();
                    }

                }
                break;
            }
        }
    }

    private TALocation linkLocation(ProcessElem toElem, TALocation prev, SequenceFlow transFlow, boolean diff)
    {
        if (diff)
        {
            // cross pool
            return linkLocationPoolDiff(toElem, prev, transFlow);
        }
        else
        {
            // same pool
            // transition  prev -> elem
            return linkLocationPoolSame(toElem, prev, transFlow);
        }
    }

    private TALocation linkLocationPoolSame(ProcessElem toElem, TALocation prev, SequenceFlow transFlow)
    {
        TATemplate t = prev.getParent();
        TALocation loc = newLocation(toElem.id, toElem.getName(), t);

        // add location
        t.addLocation(loc);
        outRefMap.put(loc.getId(), 0);

        // layout
        layoutHelper.layoutLocation(prev, loc, outRefMap.get(prev.getId()));

        // set transition
        setTransitionPoolSame(prev, loc, transFlow);
        return loc;
    }

    private TALocation linkLocationPoolDiff(ProcessElem toElem, TALocation prev, SequenceFlow transFlow)
    {
        TATemplate toT = templateRefMap.get(toElem.id);
        TALocation loc = newLocation(toElem.id, toElem.getName(), toT);

        // layout auto
        // set transition
        setTransitionPoolDiff(prev, loc, transFlow);

        // add location
        toT.addLocation(loc);

        outRefMap.put(loc.getId(), 0);
        return loc;
    }

    private TALocation linkLocationSub(ProcessElem toElem, TALocation prev, SequenceFlow transFlow)
    {
        TATemplate t = templateRefMap.get(toElem.id);
        TALocation loc = newLocation(toElem.id, toElem.getName(), t);
        nodeMap.put(loc.getId(), loc);
        TALocation ploc = newProcessLocation(toElem.id, toElem.getName(), t);
        TALocation cloc = newCompletedLocation(toElem.id, toElem.getName(), t);

        // seqFlows
        SequenceFlow l2p = new SequenceFlow(loc.getId() + "_START");
        l2p.setName(loc.getName() + "_START");
        SequenceFlow p2c = new SequenceFlow(loc.getId() + "_END");
        p2c.setName(loc.getName() + "_END");

        setTransition(prev, loc, transFlow);

        // add location
        t.addLocation(loc);
        outRefMap.put(loc.getId(), 0);
        t.addLocation(ploc);
        outRefMap.put(ploc.getId(), 0);
        t.addLocation(cloc);
        outRefMap.put(cloc.getId(), 0);

        // layout auto
        // set transitions
        setTransitionPoolSame(loc, ploc, l2p).setSync(loc.getIdName() + "_START", true);
        ta.getGlobals().addVariable("chan", loc.getIdName() + "_START");
        setTransitionPoolSame(ploc, cloc, p2c).setSync(loc.getIdName() + "_END", false);
        ta.getGlobals().addVariable("chan", loc.getIdName() + "_END");


        return cloc;
    }

    private TATransition setTransitionPoolSame(TALocation prev, TALocation to, SequenceFlow transFlow)
    {
        return setTransitionPoolSame(prev, to, transFlow, true);
    }

    private TATransition setTransitionPoolSame(TALocation prev, TALocation to, SequenceFlow transFlow, boolean layout)
    {
        return setTransitionPoolSame(prev, to, transFlow, "_TRANS", layout);
    }

    private TATransition setTransitionPoolSame(TALocation prev, TALocation to, SequenceFlow transFlow, String suffix, boolean layout)
    {
        TATemplate t = prev.getParent();

        // trans
        TATransition trans = newTransition((transFlow == null ? prev.getId() : transFlow.id) + "_TRANS", prev
                .getName(), prev.getId(), to.getId(), suffix);

        // layout
        if (layout)
        {
            layoutHelper.layoutLocation(prev, to, outRefMap.get(prev.getId()));
        }
        layoutHelper.layoutTransition(trans, prev, to);

        // add trans
        t.addTransition(trans);
        outRefMap.computeIfPresent(prev.getId(), (k, v) -> v + 1);
        return trans;
    }

    private void setTransitionPoolDiff(TALocation prev, TALocation to, SequenceFlow transFlow)
    {
        setTransitionPoolDiff(prev, to, transFlow, true);
    }

    private void setTransitionPoolDiff(TALocation prev, TALocation to, SequenceFlow transFlow, boolean layout)
    {
        TATemplate fromT = prev.getParent();
        TATemplate toT = to.getParent();

        // completed node in fromT
        TALocation completed = newSignalLocation(transFlow.id, prev.getName() + (transFlow
                .getName() == null ? "" : "_" + transFlow.getName()), fromT);

        // sync transition  prev -> completed
        TATransition prev2completed = newTransition(transFlow.id, prev.getName(), prev.getId(), completed
                .getId(), "_SEND_TRANS");

        // prev node in toT
        TALocation toPrev;
        if (stateStack.peek() == BPMNElemType.ExclusiveGateway)
            toPrev = branchStateStack.peek().get(toT.getId());
        else
            toPrev = toT.getPtr();
        assert toPrev != null;

        // sync transition  ptr -> recv
        TATransition ptr2recv = newTransition(transFlow.id, toPrev.getName(), toPrev.getId(), to
                .getId(), "_RECV_TRANS");

        // sync channels
        String channel = transFlow.id + transFlow.getName();
        prev2completed.setSync(channel, true);
        ptr2recv.setSync(channel, false);
        ta.getGlobals().addVariable("chan", channel);

        // add location
        fromT.addLocation(completed);
        outRefMap.put(completed.getId(), 0);

        // layout
        if (layout)
        {
            layoutHelper.layoutLocation(prev, completed, outRefMap.get(prev.getId()));
            layoutHelper.layoutLocation(toPrev, to, outRefMap.get(toPrev.getId()));
        }
        layoutHelper.layoutTransition(prev2completed, prev, completed);
        layoutHelper.layoutTransition(ptr2recv, toPrev, to);

        // add trans
        fromT.addTransition(prev2completed);
        outRefMap.computeIfPresent(prev.getId(), (k, v) -> v + 1);
        toT.addTransition(ptr2recv);
        outRefMap.computeIfPresent(toPrev.getId(), (k, v) -> v + 1);
    }

    private void setTransitionNoLayout(TALocation prev, TALocation to, SequenceFlow transFlow)
    {
        setTransition(prev, to, transFlow, false);
    }

    private void setTransition(TALocation prev, TALocation to, SequenceFlow transFlow)
    {
        setTransition(prev, to, transFlow, true);
    }

    private void setTransition(TALocation prev, TALocation to, SequenceFlow transFlow, boolean layout)
    {
        TATemplate fromT = prev.getParent();
        TATemplate toT = to.getParent();
        if (fromT.equals(toT))
            setTransitionPoolSame(prev, to, transFlow, layout);
        else
            setTransitionPoolDiff(prev, to, transFlow, layout);
    }

    private void errorLog(String msg)
    {
        System.out.println(msg);
        System.exit(-1);
    }

    private TALocation newLocation(String id, String name, TATemplate parent)
    {
        TALocation loc = new TALocation();
        loc.setId(ng.getIdHitFix(id)).setName(ng.getNameHitFix(name)).setParent(parent);
        return loc;
    }

    private TALocation newLocation(String id, String name, TATemplate parent, String suffix)
    {
        return newLocation(id + suffix, name + suffix, parent);
    }

    private TALocation newSignalLocation(String id, String name, TATemplate parent)
    {
        return newLocation(id, name, parent, "_SIGNAL");
    }

    private TALocation newInitLocation(String id, String name, TATemplate parent)
    {
        return newLocation(id, name, parent, "_INIT");
    }

    private TALocation newProcessLocation(String id, String name, TATemplate parent)
    {
        return newLocation(id, name, parent, "_PROCESS");
    }

    private TALocation newCompletedLocation(String id, String name, TATemplate parent)
    {
        return newLocation(id, name, parent, "_COMPLETED");
    }

    private TALocation newFinishLocation(String id, String name, TATemplate parent)
    {
        return newLocation(id, name, parent, "_FINISH");
    }

    private TATransition newTransition(String id, String name, String fromRef, String toRef)
    {
        TATransition trans = new TATransition();
        trans.setId(id).setName(name).setSourceRef(fromRef).setTargetRef(toRef);
        return trans;
    }

    private TATransition newTransition(String id, String name, String fromRef, String toRef, String suffix)
    {
        return newTransition(id + suffix, name + suffix, fromRef, toRef);
    }

    private TATemplate generateControlTemplate(Process process)
    {
        TATemplate ct = new TATemplate(process.id);
        String cName = process.getName() + "_MAIN";
        ct.setName(cName);

        // control template locations setup
        TALocation cSrt = newInitLocation(cName, cName, ct);
        TALocation cPro = newProcessLocation(cName, cName, ct);
        TALocation cEnd = newFinishLocation(cName, cName, ct);

        // transitions
        TATransition cSP = newTransition(cName, cName, cSrt.getId(), cPro.getId(), "_INIT_PROCESS");
        TATransition cPE = newTransition(cName, cName, cPro.getId(), cEnd.getId(), "_PROCESS_END");

        // sync signals
        cSP.setSync(cName + "_START", true);
        ta.getGlobals().addVariable("chan", cName + "_START");
        cPE.setSync(cName + "_END", false);
        ta.getGlobals().addVariable("chan", cName + "_END");

        // add locations
        ct.addLocation(cSrt);
        outRefMap.put(cSrt.getId(), 0);
        ct.addLocation(cPro);
        outRefMap.put(cPro.getId(), 0);
        ct.addLocation(cEnd);
        outRefMap.put(cEnd.getId(), 0);

        // layout
        layoutHelper.layoutLocation(cSrt);
        layoutHelper.layoutLocation(cSrt, cPro, outRefMap.get(cSrt.getId()));
        layoutHelper.layoutLocation(cPro, cEnd, outRefMap.get(cPro.getId()));
        layoutHelper.layoutTransition(cSP, cSrt, cPro);
        layoutHelper.layoutTransition(cPE, cPro, cEnd);

        // add trans
        ct.addTransition(cSP);
        outRefMap.computeIfPresent(cSrt.getId(), (k, v) -> v + 1);
        ct.addTransition(cPE);
        outRefMap.computeIfPresent(cPro.getId(), (k, v) -> v + 1);

        ct.setInitRef(cSrt.getId());
        return ct;
    }

    private Map<String, TATemplate> generateTemplateSet(Process process)
    {
        Map<String, TATemplate> map = new HashMap<>();
        templateRefMap.clear();
        for (var lane : process.getLaneSet().getLanes())
        {
            // template setup
            TATemplate t = new TATemplate(lane.id);
            String name = ng.getNameHitFix(process.getName() + "_" + lane.getName());
            t.setName(name);
            ta.addTemplate(t);

            // init location setup
            TALocation srt = newInitLocation(name, name, t);

            layoutHelper.reset();
            layoutHelper.layoutLocation(srt);
            outRefMap.put(srt.getId(), 0);

            t.setInitRef(srt.getId());
            t.addLocation(srt);

            for (var nodeRef : lane.getRefs())
            {
                templateRefMap.put(nodeRef, t);
            }
            t.setUnused(true);
        }

        return map;
    }
}
